from typing import Optional

from kubernetes_asyncio import client

from balder.services.kube_base import KubeBase, format_selector, format_containers_and_volumes


class JobService(KubeBase):
    async def create_job(
            self,
            namespace: str,
            name: str,
            containers: list[dict],
            volumes: Optional[list[dict]] = None,
            node_selector: Optional[dict] = None,
            init_containers: Optional[list[dict]] = None,
            retries: int = 3,
            restart_policy: str = "Never",
            completions: int = 1,
            parallelism: int = 1,
            auto_remove_pod_time: int = 259200  # 3 days
    ) -> client.V1Job:
        job_spec_template_containers, job_spec_template_volumes, job_spec_template_init_containers = format_containers_and_volumes(containers, volumes, init_containers)
        job = client.V1Job(
            api_version="batch/v1",
            kind="Job",
            metadata=client.V1ObjectMeta(
                name=name,
                labels={"app": name}
            ),
            spec=client.V1JobSpec(
                backoff_limit=retries,
                completions=completions,
                parallelism=parallelism,
                ttl_seconds_after_finished=auto_remove_pod_time,
                template=client.V1PodTemplateSpec(
                    spec=client.V1PodSpec(
                        containers=job_spec_template_containers,
                        volumes=job_spec_template_volumes,
                        init_containers=job_spec_template_init_containers,
                        node_selector=node_selector,
                        restart_policy=restart_policy
                    )
                )
            )
        )

        return await self.batch_v1_api.create_namespaced_job(namespace=namespace, body=job)

    async def delete_job(self, namespace: str, name: str) -> client.V1Status:
        return await self.batch_v1_api.delete_namespaced_job(name=name, namespace=namespace)

    async def get_job(self, namespace: str, name: str) -> client.V1Job:
        return await self.batch_v1_api.read_namespaced_job(namespace=namespace, name=name)

    async def get_job_list(self, namespace: str, label_selector: Optional[dict] = None, field_selector: Optional[dict] = None) -> client.V1JobList:
        job_label_selector, job_field_selector = format_selector(label_selector, field_selector)
        return await self.batch_v1_api.list_namespaced_job(namespace=namespace, label_selector=job_label_selector, field_selector=job_field_selector)
